import collections
import cPickle

from histogram import Histogram1D, Histogram2D, Histogram3D
from histogram.exceptions import HistogramIOError

def dump(h, path):    
    if not path.endswith('hpkl'):
        path += '.hpkl'

    if type(h) == Histogram1D or type(h) == Histogram2D or type(h) == Histogram3D:
        h = {h.name: h}
    elif isinstance(h, collections.Mapping):
        pass
    elif isinstance(h, collections.Sequence):
        h = dict([(ih.name, ih) for ih in h])
    else:
        msg = "histogram.dump() does not support output of type {0}".format(type(h))
        raise HistogramIOError(msg)

    with open(path, 'w') as fpout:
        cPickle.dump(h, fpout)
    
    return path



def load(path, names=None):
    with open(path, 'rb') as fpin:
        hd = cPickle.load(fpin)

        
    if names is not None:
        if type(names) == type('string'):
            try:
                return hd[names]
            except KeyError:
                raise HistogramIOError("Histogram `{0}' not found in file `{1}'".format(names, path))
            
        elif len(names) == 1:
            try:
                return hd[names[0]]
            except KeyError:
                raise HistogramIOError("Histogram `{0}' not found in file `{1}'".format(names, path))            
        else:
            try:
                return dict([(name, hd[name]) for name in hd])            
            except KeyError:
                raise HistogramIOError("Histogram `{0}' not found in file `{1}'".format(name, path))                
    else:    
        return hd
